#include "client.h"

EchoClient::EchoClient(const char* ip, int port, int numConnections, int maxEchoCount) 
    : ip_(ip), port_(port), numConnections_(numConnections), 
      epollFd_(-1), maxEchoCount_(maxEchoCount) {
    std::random_device rd;
    rng_.seed(rd());
}

EchoClient::~EchoClient() {
    if (epollFd_ != -1) close(epollFd_);
}

void EchoClient::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags != -1) fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

bool EchoClient::checkConnectSuccess(int fd) {
    int error = 0;
    socklen_t len = sizeof(error);
    return getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len) == 0 && error == 0;
}

int EchoClient::recvAll(int fd, char* buffer, int expectedLen) {
    int total = 0;
    int retryCount = 0;
    while (total < expectedLen && retryCount < 5) {
        int n = read(fd, buffer + total, expectedLen - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                usleep(10000);
                retryCount++;
                continue;
            }
            return -1;
        }
        if (n == 0) return total;
        total += n;
        retryCount = 0;
    }
    return total;
}

std::string EchoClient::generateTestData(int connId, int echoCount) {
    std::uniform_int_distribution<int> lenDist(5, 15);
    int dataLen = lenDist(rng_);
    std::string data;
    data += "conn" + std::to_string(connId) + "_echo" + std::to_string(echoCount) + "_";
    
    std::uniform_int_distribution<int> charDist(0, charSet_.size() - 1);
    for (int i = data.size(); i < dataLen; ++i) {
        data += charSet_[charDist(rng_)];
    }
    return data;
}

bool EchoClient::init() {
    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        return false;
    }
    return true;
}

int EchoClient::connectToServer(int connId) {
    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd == -1) return -1;

    setNonBlocking(sockfd);

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port_);
    inet_pton(AF_INET, ip_.c_str(), &addr.sin_addr);

    connect(sockfd, (struct sockaddr*)&addr, sizeof(addr));

    // 初始化连接状态：(未发送数据, 交互次数0)
    connState_[sockfd] = std::make_pair("", 0);

    struct epoll_event ev;
    ev.events = EPOLLOUT | EPOLLIN | EPOLLERR;
    ev.data.fd = sockfd;
    if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, sockfd, &ev) == -1) {
        close(sockfd);
        connState_.erase(sockfd);
        return -1;
    }

    return sockfd;
}

void EchoClient::createConnections() {
    for (int i = 0; i < numConnections_; ++i) {
        int fd = connectToServer(i);
        if (fd != -1) {
            std::cout << "Created connection " << (i+1) << "/" << numConnections_ 
                      << " (fd: " << fd << ")" << std::endl;
        } else {
            std::cerr << "Failed to create connection " << (i+1) << std::endl;
        }
        usleep(1000); // 控制连接创建速度
    }
}

void EchoClient::eventLoop() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];
    std::atomic<int> failedConnections_{0};
    std::unordered_map<int, std::pair<std::string, int>> recvBuffer;

    while (true) {
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, 1000);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            break;
        }

        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            if (event & EPOLLERR) { // 错误处理
                std::cerr << "Connection error (fd: " << fd << ")" << std::endl;
                epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                close(fd);
                failedConnections_++;
                connState_.erase(fd);
                recvBuffer.erase(fd);
                continue;
            }

            if (event & EPOLLOUT) { // 连接完成或可发送数据
                auto stateIt = connState_.find(fd);
                if (stateIt == connState_.end()) continue;

                int& echoCount = stateIt->second.second;

                // 首次连接：验证连接是否成功
                if (echoCount == 0) {
                    if (!checkConnectSuccess(fd)) {
                        std::cerr << "Connection failed (fd: " << fd << ")" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                        continue;
                    }

                    connected_++;
                    std::cout << "Connected (" << connected_ << "/" << numConnections_ 
                              << ") (fd: " << fd << ")" << std::endl;
                }

                // 生成当前交互的测试数据（connID从fd的状态中推导）
                // 实际应用中可通过connState存储connID，这里简化处理
                std::string testData = generateTestData(fd % numConnections_, echoCount + 1);
                stateIt->second.first = testData; // 更新发送的数据
                int dataLen = testData.size();

                // 发送数据
                int nwrite = write(fd, testData.c_str(), dataLen);
                if (nwrite == -1) {
                    if (errno != EAGAIN && errno != EWOULDBLOCK) {
                        std::cerr << "Send failed (fd: " << fd << ")" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                    }
                    continue;
                }

                std::cout << "Client sent (fd: " << fd 
                          << ", echo " << (echoCount + 1) << "/" << maxEchoCount_ 
                          << "): '" << testData << "' (" << nwrite << " bytes)" << std::endl;

                // 改为监听读事件，等待回显
                struct epoll_event ev;
                ev.events = EPOLLIN | EPOLLERR;
                ev.data.fd = fd;
                epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
            }

            if (event & EPOLLIN) { // 接收回显
                auto stateIt = connState_.find(fd);
                if (stateIt == connState_.end()) continue;

                std::string& expectedData = stateIt->second.first;
                int& echoCount = stateIt->second.second;
                int expectedLen = expectedData.size();

                // 初始化接收缓存
                if (recvBuffer.find(fd) == recvBuffer.end()) {
                    recvBuffer[fd] = std::make_pair("", 0);
                }
                std::string& data = recvBuffer[fd].first;
                int& received = recvBuffer[fd].second;

                // 接收数据
                int nbytes = recvAll(fd, buffer, expectedLen - received);
                if (nbytes == -1) {
                    std::cerr << "Receive failed (fd: " << fd << ")" << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    connState_.erase(fd);
                    recvBuffer.erase(fd);
                    continue;
                }

                received += nbytes;
                data.append(buffer, nbytes);
                std::cout << "Client received (fd: " << fd 
                          << "): '" << data << "' (" << received << "/" << expectedLen << " bytes)" << std::endl;

                // 验证回显
                if (received == expectedLen) {
                    if (data == expectedData) {
                        echoCount++; // 递增交互次数
                        std::cout << "Echo success (fd: " << fd 
                                  << ", completed " << echoCount << "/" << maxEchoCount_ << ")" << std::endl;

                        // 检查是否达到最大交互次数
                        if (echoCount >= maxEchoCount_) {
                            std::cout << "Max echoes reached (fd: " << fd << "), closing connection" << std::endl;
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                            connState_.erase(fd);
                            completed_++;
                        } else {
                            // 继续发送下一次数据（切换为监听写事件）
                            struct epoll_event ev;
                            ev.events = EPOLLOUT | EPOLLERR;
                            ev.data.fd = fd;
                            epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        }
                    } else {
                        std::cerr << "Echo mismatch (fd: " << fd 
                                  << "): expected '" << expectedData << "', got '" << data << "'" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                    }
                    recvBuffer.erase(fd);
                } else if (received > expectedLen) {
                    std::cerr << "Invalid length (fd: " << fd << ")" << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    connState_.erase(fd);
                    recvBuffer.erase(fd);
                }
            }
        }

        // 所有连接完成或失败时退出
        if (completed_ + failedConnections_ >= numConnections_) {
            break;
        }
    }

    std::cout << "\n===== Test Result =====" << std::endl;
    std::cout << "Total connections: " << numConnections_ << std::endl;
    std::cout << "Successfully connected: " << connected_ << std::endl;
    std::cout << "Completed all echoes: " << completed_ << std::endl;
    std::cout << "Failed connections: " << failedConnections_ << std::endl;
}

void EchoClient::start() {
    std::thread connThread(&EchoClient::createConnections, this);
    eventLoop();
    connThread.join();
}
